package st.learnopengl.apps.assets

import glm_.mat4x4.Mat4
import glm_.vec3.Vec3
import gln.BufferTarget
import gln.Usage
import gln.VertexAttrType
import gln.glf.VertexAttribute
import gln.identifiers.GlBuffer
import gln.identifiers.GlProgram
import kool.Stack
import st.learnopengl.common.GLUT
import org.lwjgl.assimp.*
import org.lwjgl.opengl.GL11C
import org.lwjgl.opengl.GL15C.*
import java.lang.RuntimeException

class Mesh(private val mesh: AIMesh, val model: st.learnopengl.apps.assets.Model) {

    private val pvb = 16
    private val boneIndexWeightStride = pvb * 2

    private var materialIndex: Int = -1

    var vertexBuffer: GlBuffer? = null
    var uvBuffer: GlBuffer? = null
    var normalBuffer: GlBuffer? = null
    var indexBuffer: GlBuffer? = null
    var bonesBuffer: GlBuffer? = null

    val elementCount = mesh.mNumFaces() * 3
    val aPos: VertexAttribute = VertexAttribute(0, 3, VertexAttrType.FLOAT, false, 3 * 4, 0)
    val aUV: VertexAttribute = VertexAttribute(1, 2, VertexAttrType.FLOAT, false, 2 * 4, 0)

    val aBoneId1: VertexAttribute = VertexAttribute(2, 4, VertexAttrType.FLOAT, false, 4 * pvb * 2, 4 * 0 * 4)
    val aBoneId2: VertexAttribute = VertexAttribute(3, 4, VertexAttrType.FLOAT, false, 4 * pvb * 2, 4 * 1 * 4)
    val aBoneId3: VertexAttribute = VertexAttribute(4, 4, VertexAttrType.FLOAT, false, 4 * pvb * 2, 4 * 2 * 4)
    val aBoneId4: VertexAttribute = VertexAttribute(5, 4, VertexAttrType.FLOAT, false, 4 * pvb * 2, 4 * 3 * 4)

    val aBoneWeight1: VertexAttribute = VertexAttribute(6, 4, VertexAttrType.FLOAT, false, 4 * pvb * 2, 4 * 4 * 4)
    val aBoneWeight2: VertexAttribute = VertexAttribute(7, 4, VertexAttrType.FLOAT, false, 4 * pvb * 2, 4 * 5 * 4)
    val aBoneWeight3: VertexAttribute = VertexAttribute(8, 4, VertexAttrType.FLOAT, false, 4 * pvb * 2, 4 * 6 * 4)
    val aBoneWeight4: VertexAttribute = VertexAttribute(9, 4, VertexAttrType.FLOAT, false, 4 * pvb * 2, 4 * 7 * 4)

    val boneMapping: HashMap<String, Int> = HashMap()
    var bones: ArrayList<Bone> = ArrayList()

    private var mGlobalInverseTransform: Mat4 = Mat4(1.0f)

    fun build() {

        materialIndex = mesh.mMaterialIndex()

        mGlobalInverseTransform = GLUT.aiMatrixToMat4(model.scene.mRootNode()!!.mTransformation())
        mGlobalInverseTransform.inverseAssign()

        vertexBuffer = GlBuffer.gen()
        vertexBuffer!!.bound(BufferTarget.ARRAY) {
            val aivector = mesh.mVertices()

            nglBufferData(
                GL_ARRAY_BUFFER,
                (AIVector3D.SIZEOF * aivector.remaining()).toLong(),
                aivector.address(),
                GL_STATIC_DRAW)
        }

        uvBuffer = GlBuffer.gen()
        uvBuffer!!.bound(BufferTarget.ARRAY) {
            Stack {
                val uvFloatBuffer = it.callocFloat(mesh.mNumVertices() * 2)


                for (i in 0 until mesh.mNumVertices()) {
                    val vector3D = mesh.mTextureCoords(0)!!.get(i)
                    uvFloatBuffer.put(vector3D.x())
                    uvFloatBuffer.put(vector3D.y())
                }

                uvFloatBuffer.flip()

                data(uvFloatBuffer)
            }

        }

        val normVector = mesh.mNormals()
        if (null != normVector) {
            normalBuffer = GlBuffer.gen()
            normalBuffer!!.bound(BufferTarget.ARRAY) {
                nglBufferData(
                    GL_ARRAY_BUFFER,
                    (AIVector3D.SIZEOF * normVector.remaining()).toLong(),
                    normVector.address(),
                    GL_STATIC_DRAW
                )
            }
        }

        if (mesh.mNumBones() > 0) {

            val boneIndexWeight = FloatArray(boneIndexWeightStride * mesh.mNumVertices()) { 0.0f }

            loadBones(boneIndexWeight, boneIndexWeightStride, pvb, pvb)

            bonesBuffer = GlBuffer.gen()
            bonesBuffer!!.bound(BufferTarget.ARRAY) {
                data(boneIndexWeight)
            }
        }

        Stack {
            val elementCount: Int = mesh.mNumFaces() * 3
            val indexDataBuffer = it.callocInt(elementCount)

            mesh.mFaces().forEach {
                if (it.mNumIndices() != 3) {
                    throw RuntimeException("AIFace.mNumIndeces is not 3 bu ${it.mNumIndices()}")
                }

                indexDataBuffer.put(it.mIndices())
            }
            indexDataBuffer.flip()

            indexBuffer = GlBuffer.gen()
            indexBuffer!!.bound(BufferTarget.ELEMENT_ARRAY) {
                data(indexDataBuffer, Usage.STATIC_DRAW)
            }
        }
    }


    fun render(program: GlProgram, runningTime: Float) {



        updateBoneTransform(runningTime)

        model.materials[materialIndex].bind(program)

        val uniformGBones = Array<Mat4>(bones.size) {
            bones[it].finalTransform
        }
        program.programUniform(program.getUniformLocation("gBones"), uniformGBones)


        vertexBuffer!!.bind(BufferTarget.ARRAY)
        aPos.enable()
        aPos.pointer()

        uvBuffer!!.bind(BufferTarget.ARRAY)
        aUV.enable()
        aUV.pointer()


        bonesBuffer!!.bind(BufferTarget.ARRAY)

        aBoneId1.enable()
        aBoneId1.pointer()

        aBoneId2.enable()
        aBoneId2.pointer()

        aBoneId3.enable()
        aBoneId3.pointer()

        aBoneId4.enable()
        aBoneId4.pointer()

        aBoneWeight1.enable()
        aBoneWeight1.pointer()

        aBoneWeight2.enable()
        aBoneWeight2.pointer()

        aBoneWeight3.enable()
        aBoneWeight3.pointer()

        aBoneWeight4.enable()
        aBoneWeight4.pointer()

        indexBuffer!!.bind(BufferTarget.ELEMENT_ARRAY)
        GL11C.glDrawElements(GL11C.GL_TRIANGLES, elementCount, GL11C.GL_UNSIGNED_INT, 0)


        vertexBuffer!!.unbind(BufferTarget.ARRAY)
        uvBuffer!!.unbind(BufferTarget.ARRAY)
        bonesBuffer!!.unbind(BufferTarget.ARRAY)
        indexBuffer!!.unbind(BufferTarget.ELEMENT_ARRAY)

        aPos.disable()
        aUV.disable()

        aBoneId1.disable()
        aBoneId2.disable()
        aBoneId3.disable()
        aBoneId4.disable()

        aBoneWeight1.disable()
        aBoneWeight2.disable()
        aBoneWeight3.disable()
        aBoneWeight4.disable()
    }

    /**
     * helper function for load bone info from AiMesh
     */
    private fun loadBones(boneIndexWeight: FloatArray, stride: Int, indexSize: Int, weightSize: Int) {
        for (i in 0 until mesh.mNumBones()) {
            var boneIndex = 0
            val aiBone = AIBone.create(mesh.mBones()!!.get(i))

            val boneName = aiBone.mName().dataString()

            if (boneMapping.containsKey(boneName)) {
                boneIndex = boneMapping.get(boneName)!!
            } else {
                boneIndex = boneMapping.size
                bones.add(Bone())
            }

            boneMapping[boneName] = boneIndex
            bones[boneIndex].setAiMatrix(aiBone.mOffsetMatrix())

            for (j in 0 until aiBone.mNumWeights()) {
                val vertexId = aiBone.mWeights().get(j).mVertexId()
                val weight = aiBone.mWeights().get(j).mWeight()
                addBoneData(boneIndexWeight, boneIndex, weight, vertexId * stride, indexSize, weightSize)
            }
        }
    }


    private fun addBoneData(boneIndexWeight: FloatArray, boneID: Int, boneWeight: Float, vertexIndex: Int, indexSize: Int, weightSize: Int) {
        if (indexSize != weightSize) {
            throw RuntimeException("indexSize[$indexSize] not equal weightSize [$weightSize]")
        }

        for (i in 0 until weightSize) {
            if (boneIndexWeight[vertexIndex + indexSize + i] == 0.0f) {
                boneIndexWeight[vertexIndex + i] = boneID.toFloat()
                boneIndexWeight[vertexIndex + indexSize + i] = boneWeight;
                return
            }
        }

        throw RuntimeException("bone index out of range")
    }


    private fun readNodeHeirarchy(animationTime: Float, aiNode: AINode, parentTransform: Mat4) {
        val nodeName = aiNode.mName().dataString()

        val nodeTransformation = aiNode.mTransformation()

        var nodeMat4 = GLUT.aiMatrixToMat4(nodeTransformation)
        val aiAnimation = AIAnimation.create(model.scene.mAnimations()!!.get(0))

        val aiNodeAnim = findNodeAnim(aiAnimation, nodeName)
        if (aiNodeAnim != null) {
            val scaling = calculateInterpolatedScaling(animationTime, aiNodeAnim)
            val scalingMatrix = Mat4(1.0)
            scalingMatrix.scale(scaling)

            val rotationMatrix = calculateInterpolatedRotation(animationTime, aiNodeAnim)

            val translateVec3 = calculateInterpolatedPosition(animationTime, aiNodeAnim)

            val translateMat4 = Mat4().translate(translateVec3)

            nodeMat4 = translateMat4 * rotationMatrix * scalingMatrix
        }

        val globalTransform = parentTransform * nodeMat4

        if (boneMapping.containsKey(nodeName)) {
            val boneIndex = boneMapping[nodeName]!!
            bones[boneIndex].finalTransform = mGlobalInverseTransform * globalTransform * bones[boneIndex].matrix
        }

        for (i in 0 until aiNode.mNumChildren()) {
            readNodeHeirarchy(animationTime, AINode.create(aiNode.mChildren()!!.get(i)), globalTransform)
        }
    }

    fun updateBoneTransform(timeInSeconds: Float) {
        val m = Mat4()
        m.identity()
        readNodeHeirarchy(timeInSeconds, model.scene.mRootNode()!!, m)
    }

    fun findNodeAnim(aiAnimation: AIAnimation, nodeName: String): AINodeAnim? {
        for (i in 0 until aiAnimation.mNumChannels()) {
            val aiNodeAnimAddr = aiAnimation.mChannels()!!.get(i)
            val aiNodeAnim = AINodeAnim.create(aiNodeAnimAddr)
            if (aiNodeAnim.mNodeName().dataString() == nodeName) {
                return aiNodeAnim
            }
        }

        return null
    }

    fun findScaling(animationTime: Float, aiNodeAnim: AINodeAnim): Int {
        assert (aiNodeAnim.mNumScalingKeys() > 0)

        for (i in 0 until aiNodeAnim.mNumScalingKeys() - 1) {
            if (animationTime < aiNodeAnim.mScalingKeys()!!.get(i + 1).mTime().toFloat()) {
                return i
            }
        }
        throw RuntimeException("time out of bounds $animationTime")
        return 0
    }

    fun findRotation(animationTime: Float, aiNodeAnim: AINodeAnim): Int {
        assert(aiNodeAnim.mNumRotationKeys() > 0)

        for (i in 0 until aiNodeAnim.mNumRotationKeys() - 1) {
            if (animationTime < aiNodeAnim.mRotationKeys()!!.get(i + 1).mTime().toFloat()) {
                return i
            }
        }

        assert(false)

        return 0
    }

    fun findPosition(animationTime: Float, aiNodeAnim: AINodeAnim): Int {
        assert(aiNodeAnim.mNumPositionKeys() > 0)

        for (i in 0 until aiNodeAnim.mNumPositionKeys() - 1) {
            if (animationTime < aiNodeAnim.mPositionKeys()!!.get(i + 1).mTime().toFloat()) {
                return i
            }
        }

        assert(false)
        return 0
    }

    fun calculateInterpolatedScaling(animationTime: Float, aiNodeAnim: AINodeAnim): Vec3 {
        val ret = Vec3()

        if (aiNodeAnim.mNumScalingKeys() == 1) {
            val aiVec3 = aiNodeAnim.mScalingKeys()!!.get(0).mValue()
            ret.x = aiVec3.x()
            ret.y = aiVec3.y()
            ret.z = aiVec3.z()
            return ret
        }

        val scalingIndex = findScaling(animationTime, aiNodeAnim)
        val nextScalingIndex = scalingIndex + 1

        assert(nextScalingIndex < aiNodeAnim.mNumScalingKeys())
        val deltaTime = (aiNodeAnim.mScalingKeys()!!.get(nextScalingIndex).mTime() - aiNodeAnim.mScalingKeys()!!.get(scalingIndex).mTime()).toFloat()
        var factor = (animationTime - aiNodeAnim.mScalingKeys()!!.get(scalingIndex).mTime()).toFloat() / deltaTime

        if (factor.isInfinite() || factor.isNaN()) {
            factor = 0.0f
        }

        val start = GLUT.aiVector3DToVec3(aiNodeAnim.mScalingKeys()!!.get(scalingIndex).mValue())
        val end = GLUT.aiVector3DToVec3(aiNodeAnim.mScalingKeys()!!.get(nextScalingIndex).mValue())

        val deltaVector = end - start

        return start + deltaVector * factor
    }

    fun calculateInterpolatedRotation(animationTime: Float, aiNodeAnim: AINodeAnim): Mat4 {
        if (aiNodeAnim.mNumRotationKeys() == 1) {
            val aiQuaternion = aiNodeAnim.mRotationKeys()!!.get(0).mValue()

            val quat = GLUT.aiQuaternionToQuat(aiQuaternion)

            return quat.toMat4()
        }

        val rotationIndex = findRotation(animationTime, aiNodeAnim)
        val nextRotationIndex = rotationIndex + 1
        assert(nextRotationIndex < aiNodeAnim.mNumRotationKeys())

        val deltaTime = (aiNodeAnim.mRotationKeys()!!.get(nextRotationIndex).mTime() - aiNodeAnim.mRotationKeys()!!.get(rotationIndex).mTime()).toFloat()
        var factor = (animationTime - aiNodeAnim.mRotationKeys()!!.get(rotationIndex).mTime().toFloat()) / deltaTime

        assert(factor in 0.0f..1.0f)

        if (factor.isInfinite() || factor.isNaN()) {
            factor = 0.0f
        }

        val start = GLUT.aiQuaternionToQuat(aiNodeAnim.mRotationKeys()!!.get(rotationIndex).mValue())
        val end = GLUT.aiQuaternionToQuat(aiNodeAnim.mRotationKeys()!!.get(nextRotationIndex).mValue())
        val deltaQuat = end - start


        val ret = GLUT.interpolateQuate(start, end, factor)
        ret.normalizeAssign()

        return ret.toMat4()
    }

    fun calculateInterpolatedPosition(animationTime: Float, aiNodeAnim: AINodeAnim): Vec3 {
        if (aiNodeAnim.mNumPositionKeys() == 1) {
            return GLUT.aiVector3DToVec3(aiNodeAnim.mPositionKeys()!!.get(0).mValue())
        }

        val positionIndex = findPosition(animationTime, aiNodeAnim)
        val nextPositionIndex = positionIndex + 1
        assert(nextPositionIndex < aiNodeAnim.mNumPositionKeys())

        val deltaTime = (aiNodeAnim.mPositionKeys()!!.get(nextPositionIndex).mTime() - aiNodeAnim.mPositionKeys()!!.get(nextPositionIndex).mTime()).toFloat()
        var factor = (animationTime - aiNodeAnim.mPositionKeys()!!.get(positionIndex).mTime().toFloat()) / deltaTime

        assert(factor in 0.0f .. 1.0f)

        if (factor.isInfinite() || factor.isNaN()) {
            factor = 0.0f
        }

        val start = GLUT.aiVector3DToVec3(aiNodeAnim.mPositionKeys()!!.get(positionIndex).mValue())
        val end = GLUT.aiVector3DToVec3(aiNodeAnim.mPositionKeys()!!.get(nextPositionIndex).mValue())
        val delta = end - start

        return start + delta * factor
    }
}